import tensorflow as tf
import cv2
import numpy as np
import os

# 设置GPU内存按需增长
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

# 定义待测图片路径，您需要将其修改为您自己的路径
img_path = r"/data/wengjj/catdog/data/cats"


def traverse_directory(directory):
    for root, dirs, files in os.walk(directory):
        for file in files:
            print(os.path.join(root, file))
            i = os.path.join(root, file)
            image = cv2.imread(i)
            image = cv2.resize(image, (img_height, img_width))
            # 将图片转换为数组
            image = np.expand_dims(image, axis=0)
            image = image / 255.0
            # 使用模型进行预测
            predictions = model.predict(image)
            print(predictions)
            if predictions[0] < 0.5:
                # correct_num += 1
                print(f'{i}.jpg: 狗  预测正确')
            elif predictions[0] >= 0.5:
                # correct_num += 1
                print(f'{i}.jpg: 猫  预测正确')
            else:
                if predictions[0] < 0.5:
                    print(f'{i}.jpg: 狗  预测错误')
                else:
                    print(f'{i}.jpg: 猫  预测错误')


# 加载模型
model = tf.keras.models.load_model(r'/data/wengjj/catdog/cat_dog_classifier30.h5')

# 定义图像尺寸
img_height, img_width = 150, 150

traverse_directory(img_path)
